// https://leetcode.cn/problems/binary-tree-pruning/description/

// 算法思路总结：
// 1. 后序遍历递归剪枝二叉树
// 2. 先递归处理左右子树，再处理当前节点
// 3. 如果当前节点是值为0的叶子节点，返回nullptr（剪枝）
// 4. 否则保留当前节点，返回修剪后的子树
// 5. 时间复杂度：O(n)，空间复杂度：O(h)（递归栈深度）

#include <iostream>
using namespace std;

#include <vector>
#include <algorithm>
#include "BinaryTreeUtils.h"

class Solution 
{
public:
    TreeNode* pruneTree(TreeNode* root) 
    {
        if (root == nullptr) return nullptr;

        root->left = pruneTree(root->left);
        root->right = pruneTree(root->right);

        if (root->left == nullptr && root->right == nullptr && root->val == 0) 
            return nullptr;

        return root; 
    }
};
int main()
{
    vector<string> tree1 = {"1", "null", "0", "0", "1"};
    vector<string> tree2 = {"1", "0", "1", "0", "0", "0", "1"};
    Solution sol;

    auto root1 = buildTree(tree1);
    auto root2 = buildTree(tree2);

    auto r1 = sol.pruneTree(root1);
    auto r2 = sol.pruneTree(root2);

    printTree(r1);
    printTree(r2);

    return 0;
}